package mongo

import (
	"bytes"
	"context"
	"strings"
	"time"

	mongoProm "github.com/globocom/mongo-go-prometheus"
	"github.com/pkg/errors"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/gridfs"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readpref"
	"go.mongodb.org/mongo-driver/x/bsonx"
)

type TraverseFunc func(data interface{}) error

type Index struct {
	Ver    int              `bson:"v"`
	Name   string           `bson:"name"`
	Key    map[string]int32 `bson:"key"`
	Unique bool             `bson:"unique"`
}

var _ Client = (*client)(nil)

type Client interface {
	InsertOne(ctx context.Context, table string, data interface{}) error
	InsertMany(ctx context.Context, table string, data []interface{}) error
	FindOne(ctx context.Context, table string, filter, data interface{}) error
	FindOneWithProjection(ctx context.Context, table string, filter, projection, data interface{}) error
	Find(ctx context.Context, table string, filter, data interface{}) error
	FindWithProjection(ctx context.Context, table string, filter, projection, data interface{}) error
	UpdateOne(ctx context.Context, table string, filter, data interface{}) error
	UpsertOne(ctx context.Context, table string, filter, data interface{}) error
	Update(ctx context.Context, table string, filter interface{}, data interface{}) error
	FindOneAndUpdate(ctx context.Context, table string, filter, update, data interface{}) error
	DeleteOne(ctx context.Context, table string, filter interface{}) error
	DeleteAll(ctx context.Context, table string, filter interface{}) (int64, error)

	MultiReplaceInsert(ctx context.Context, table string, filter []interface{}, data []interface{}) error
	RunJavascript(ctx context.Context, script string) ([]interface{}, error)
	Traverse(ctx context.Context, table string, finder, data, projection interface{}, limit int64, fun TraverseFunc) error
	Transaction(ctx context.Context, table string) error
	Session(ctx context.Context, table string) error
	SimpleTransaction(ctx context.Context, f func(sessionContext context.Context) error) error

	GetIndexes(ctx context.Context, table string) ([]Index, error)
	CreateOneIndex(ctx context.Context, table string, unique bool, keys interface{}) (string, error)
	Count(ctx context.Context, table string, finder interface{}) (count int64, err error)

	LoadFileData(filename string, filepath string) ([]byte, error)
	SaveFileData(fileName string, filePath string, fileData []byte) error
	DeleteFileData(fileName string, filePath string) error

	InsertOneResult(ctx context.Context, table string, data interface{}) (*mongo.InsertOneResult, error)
	UpdateOneResult(ctx context.Context, table string, filter interface{}, data interface{}) (*mongo.UpdateResult, error)
	UpdateOneOptions(ctx context.Context, table string, filter, data interface{}, opts ...*options.UpdateOptions) error
	FindOptions(ctx context.Context, table string, finder interface{}, data interface{}, opts ...*options.FindOptions) error
	FindOneOptions(ctx context.Context, table string, finder interface{}, data interface{}, opts ...*options.FindOneOptions) error
}

type Config struct {
	Namespace   string
	Hosts       []string
	Database    string
	UserName    string
	Password    string
	AuthSource  string
	MaxPoolSize uint
	MinPoolSize uint
	MaxIdleTime uint
}

func NewClient(conf *Config) (Client, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)

	opts := options.Client().SetMaxPoolSize(uint64(conf.MaxPoolSize)).SetMinPoolSize(uint64(conf.MinPoolSize)).
		SetMaxConnIdleTime(time.Duration(conf.MaxIdleTime) * time.Second).SetHosts(conf.Hosts).
		SetSocketTimeout(3 * time.Second).SetServerSelectionTimeout(2 * time.Second).SetMonitor(mongoProm.NewCommandMonitor(
		mongoProm.WithInstanceName(conf.Database),
		mongoProm.WithNamespace(conf.Namespace),
		mongoProm.WithDurationBuckets([]float64{.001, .005, .01, 0.1, 0.5, 1, 2, 3, 4, 5, 5.5, 6, 7, 8, 10, 11}),
	))

	if conf.UserName != "" && conf.Password != "" {
		opts = opts.SetAuth(options.Credential{
			Username:   conf.UserName,
			Password:   conf.Password,
			AuthSource: "admin",
		})
	}

	cli, err := mongo.Connect(ctx, opts)
	defer cancel()

	if err != nil {
		return nil, err
	}

	// 判断服务是否可用
	if err = cli.Ping(ctx, readpref.Primary()); err != nil {
		return nil, err
	}

	return &client{
		cli:  cli,
		conf: conf,
	}, nil
}

type client struct {
	cli  *mongo.Client
	conf *Config
}

func (c *client) InsertOne(ctx context.Context, table string, data interface{}) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	_, err := collection.InsertOne(ctx, data)
	return err
}

func (c *client) InsertMany(ctx context.Context, table string, data []interface{}) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	_, err := collection.InsertMany(ctx, data)
	return err
}

func (c *client) FindOne(ctx context.Context, table string, filter, data interface{}) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	return collection.FindOne(ctx, filter).Decode(data)
}

func (c *client) FindOneWithProjection(ctx context.Context, table string, filter interface{},
	projection interface{}, data interface{}) error {
	opt := options.FindOne().SetProjection(projection)
	collection := c.cli.Database(c.conf.Database).Collection(table)
	return collection.FindOne(ctx, filter, opt).Decode(data)
}

func (c *client) Find(ctx context.Context, table string, filter, data interface{}) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	cursor, err := collection.Find(ctx, filter)
	if err != nil {
		return err
	}

	return cursor.All(ctx, data)
}

func (c *client) FindWithProjection(ctx context.Context, table string, filter, projection, data interface{}) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	opt := options.Find().SetProjection(projection)
	cursor, err := collection.Find(ctx, filter, opt)
	if err != nil {
		return err
	}

	return cursor.All(ctx, data)
}

func (c *client) Update(ctx context.Context, table string, filter interface{}, data interface{}) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	_, err := collection.UpdateMany(ctx, filter, data)
	return err
}

func (c *client) UpdateOne(ctx context.Context, table string, filter, data interface{}) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	_, err := collection.UpdateOne(ctx, filter, data)
	if err != nil {
		return err
	}

	return nil
}

func (c *client) UpsertOne(ctx context.Context, table string, filter, data interface{}) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	opts := options.Update().SetUpsert(true)
	_, err := collection.UpdateOne(ctx, filter, data, opts)
	if err != nil {
		return err
	}

	return nil
}

func (c *client) FindOneAndUpdate(ctx context.Context, table string, filter, update, data interface{}) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	opts := options.FindOneAndUpdate().SetReturnDocument(options.After)
	err := collection.FindOneAndUpdate(ctx, filter, update, opts).Decode(data)
	if err != nil {
		return err
	}

	return nil
}

func (c *client) DeleteOne(ctx context.Context, table string, filter interface{}) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	_, err := collection.DeleteOne(ctx, filter)
	if err != nil {
		return err
	}

	return nil
}

func (c *client) DeleteAll(ctx context.Context, table string, filter interface{}) (int64, error) {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	result, err := collection.DeleteMany(ctx, filter)
	if err != nil {
		return 0, err
	}

	return result.DeletedCount, nil
}

func (c *client) Traverse(ctx context.Context, table string, filter, data, projection interface{}, limit int64, fun TraverseFunc) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	batchSize := int32(200)
	skipCount := int64(0)

	cursor, err := collection.Find(ctx, filter, &options.FindOptions{
		BatchSize:  &batchSize,
		Skip:       &skipCount,
		Projection: projection,
		Limit:      &limit,
	})

	if err != nil {
		return err
	}

	for cursor.TryNext(context.TODO()) {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			err = cursor.Decode(data)
			if err != nil {
				return err
			}

			err := fun(data)
			if err != nil {
				return err
			}

			length := cursor.RemainingBatchLength()
			if length == 0 {
				time.Sleep(time.Millisecond * 100)
			}
		}
	}

	return nil
}

// MultiReplaceInsert ...
func (c *client) MultiReplaceInsert(ctx context.Context, table string, filter, data []interface{}) error {
	length := len(filter)
	if length != len(data) {
		return errors.New("filter and data do not match")
	}

	models := make([]mongo.WriteModel, 0, length)
	for i := 0; i < length; i++ {
		model := mongo.NewReplaceOneModel().SetFilter(filter[i])
		model.SetReplacement(data[i])
		model.SetUpsert(true)
		models = append(models, model)
	}

	collection := c.cli.Database(c.conf.Database).Collection(table)
	_, err := collection.BulkWrite(ctx, models, options.BulkWrite().SetOrdered(false))
	if err != nil {
		return err
	}

	return nil
}

// RunJavascript 需要mongodb 4.0之前的版本
func (c *client) RunJavascript(ctx context.Context, script string) ([]interface{}, error) {
	js := bsonx.JavaScript(script)
	res, err := c.cli.Database(c.conf.Database).RunCommand(ctx, bson.M{"eval": js}).DecodeBytes()
	if err != nil {
		return nil, err
	}

	items, err := res.LookupErr("retval")
	if err != nil {
		return nil, err
	}

	arr := make([]interface{}, 0)
	err = items.Unmarshal(&arr)
	if err != nil {
		return nil, err
	}

	return arr, nil
}

// Transaction ...
func (c *client) Transaction(ctx context.Context, table string) error {
	session, err := c.cli.StartSession()
	if err != nil {
		return err
	}
	defer session.EndSession(ctx)

	// session 上下文来保证一个事务
	sessCtx := mongo.NewSessionContext(ctx, session)

	// 开始事务
	err = session.StartTransaction()
	if err != nil {
		return err
	}

	// 可以用多个表，保证事务执行
	collection := session.Client().Database(c.conf.Database).Collection(table)

	_, err = collection.InsertOne(sessCtx, bson.M{"_id": "222", "name": "ddd", "age": 50})
	if err != nil {
		return err
	}

	// 写重复id
	_, err = collection.InsertOne(sessCtx, bson.M{"_id": "111", "name": "ddd", "age": 50})
	if err != nil {
		_ = session.AbortTransaction(ctx)
		return err
	}

	return session.CommitTransaction(ctx)
}

// Session 需要Replication Set才能执行
func (c *client) Session(ctx context.Context, table string) error {
	return c.cli.UseSession(ctx, func(sessionContext mongo.SessionContext) error {
		err := sessionContext.StartTransaction()
		if err != nil {
			return err
		}

		col := sessionContext.Client().Database(c.conf.Database).Collection(table)
		_, err = col.InsertOne(sessionContext, bson.M{"_id": "444", "name": "ddd", "age": 50})
		if err != nil {
			return err
		}

		_, err = col.InsertOne(sessionContext, bson.M{"_id": "111", "name": "ddd", "age": 50})
		if err != nil {
			_ = sessionContext.AbortTransaction(sessionContext)
			return err
		}

		return sessionContext.CommitTransaction(sessionContext)
	})
}

func (c *client) SimpleTransaction(ctx context.Context, f func(sessionContext context.Context) error) error {
	if f == nil {
		return errors.New("callback function is nil")
	}

	return c.cli.UseSession(ctx, func(sessionContext mongo.SessionContext) error {
		if err := sessionContext.StartTransaction(); err != nil {
			return err
		}

		if err := f(sessionContext); err != nil {
			_ = sessionContext.AbortTransaction(sessionContext)
			return err
		}

		return sessionContext.CommitTransaction(sessionContext)
	})
}

func (c *client) GetIndexes(ctx context.Context, table string) ([]Index, error) {
	collection := c.cli.Database(c.conf.Database).Collection(table)

	cursor, err := collection.Indexes().List(ctx)
	if err != nil {
		return nil, err
	}

	indexes := make([]Index, 0)
	if err = cursor.All(ctx, &indexes); err != nil {
		return nil, err
	}

	return indexes, nil
}

func (c *client) CreateOneIndex(ctx context.Context, table string, unique bool, keys interface{}) (string, error) {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	return collection.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys:    keys,
		Options: options.Index().SetUnique(unique),
	})
}

func (c *client) Count(ctx context.Context, table string, finder interface{}) (count int64, err error) {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	count, err = collection.CountDocuments(ctx, finder)
	if err != nil {
		return 0, err
	}
	return
}

// mongofiles 命令也可以操作文件
func (c *client) getGridFsBucket(filePath string) (*gridfs.Bucket, error) {
	colName := strings.ReplaceAll(filePath, "/", "_")
	opts := options.GridFSBucket().SetName(colName)
	bucket, err := gridfs.NewBucket(c.cli.Database(c.conf.Database), opts)
	return bucket, err
}

func (c *client) SaveFileData(filePath, fileName string, fileData []byte) error {
	// TODO: 校验文件MD5,避免重复存储文件
	bucket, err := c.getGridFsBucket(filePath)
	if err != nil {
		return err
	}

	return bucket.UploadFromStreamWithID(fileName, fileName, bytes.NewReader(fileData))
}

func (c *client) LoadFileData(filePath, fileName string) ([]byte, error) {
	bucket, err := c.getGridFsBucket(filePath)
	if err != nil {
		return nil, err
	}

	fileBuffer := bytes.NewBuffer(nil)
	if _, err = bucket.DownloadToStreamByName(fileName, fileBuffer); err != nil {
		return nil, err
	}

	return fileBuffer.Bytes(), nil
}

func (c *client) DeleteFileData(filePath, fileName string) error {
	bucket, err := c.getGridFsBucket(filePath)
	if err != nil {
		return err
	}

	return bucket.Delete(fileName)
}

func (c *client) InsertOneResult(ctx context.Context, table string, data interface{}) (*mongo.InsertOneResult, error) {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	result, err := collection.InsertOne(ctx, data)
	if err != nil {
		return nil, err
	}
	return result, nil
}

func (c *client) UpdateOneResult(ctx context.Context, table string, filter interface{}, data interface{}) (*mongo.UpdateResult, error) {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	res, err := collection.UpdateOne(ctx, filter, data)
	if err != nil {
		return nil, err
	}

	return res, nil
}

func (c *client) UpdateOneOptions(ctx context.Context, table string, filter interface{}, data interface{}, opts ...*options.UpdateOptions) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	_, err := collection.UpdateOne(ctx, filter, data, opts...)
	if err != nil {
		return err
	}

	return nil
}

func (c *client) FindOneOptions(ctx context.Context, table string, finder interface{}, data interface{}, opts ...*options.FindOneOptions) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	return collection.FindOne(ctx, finder, opts...).Decode(data)
}

func (c *client) FindOptions(ctx context.Context, table string, finder interface{}, data interface{}, opts ...*options.FindOptions) error {
	collection := c.cli.Database(c.conf.Database).Collection(table)
	cursor, err := collection.Find(ctx, finder, opts...)
	if err != nil {
		return err
	}

	return cursor.All(context.TODO(), data)
}
